import misc
using .misc: make_tuple

""" Numerical differentiation """

function numerical_diff(func, ssinputs_dict::Dict{String, Any}, shock_dict::Dict{String, Any}, h::Float64=1E-4, y_ss_list::Union{Nothing, Vector{Any}}=nothing)
     """Differentiate function numerically via forward difference, i.e. calculate

    f'(xss)*shock = (f(xss + h*shock) - f(xss))/h

    for small h. (Variable names inspired by application of differentiating around ss.)

    Parameters
    ----------
    func            : function, 'f' to be differentiated
    ssinputs_dict   : dict, values in 'xss' around which to differentiate
    shock_dict      : dict, values in 'shock' for which we're taking derivative
                        (keys in shock_dict are weak subset of keys in ssinputs_dict)
    h               : [optional] scalar, scaling of forward difference 'h'
    y_ss_list       : [optional] list, value of y=f(xss) if we already have it

    Returns
    ----------
    dy_list : list, output f'(xss)*shock of numerical differentiation
    """

    # compute ss output if not supplied
    if y_ss_list == nothing
        y_ss_list = make_tuple(func(ssinputs_dict...))
    end

    # response to small shock
    shocked_inputs = Dict(k => ssinputs_dict[k] + h * get(shock_dict, k, 0.0) for k in keys(ssinputs_dict))
    y_list = make_tuple(func(shocked_inputs...))

    # scale responses back up, dividing by h
    dy_list = [(y - y_ss) / h for (y, y_ss) in zip(y_list, y_ss_list)]

    return dy_list

end

# test case for numerical_diff to see if this silly little function works idk it prob won't

# using Test

# Dummy make_tuple function for testing
# make_tuple(x) = (x,)

# Define a simple test function to differentiate
# function test_func(x, y)
    # return x^2 + y^2
# end

# Define the steady-state inputs and shocks
# ssinputs_dict = Dict("x" => 1.0, "y" => 2.0)
# shock_dict = Dict("x" => 0.1, "y" => 0.2)

# Expected derivatives (analytically: df/dx = 2x, df/dy = 2y)
# expected_diff = [2.0 * ssinputs_dict["x"] * shock_dict["x"], 2.0 * ssinputs_dict["y"] * shock_dict["y"]]

# Testing numerical_diff
# dy_list = numerical_diff(test_func, ssinputs_dict, shock_dict)
# @test isapprox(dy_list, expected_diff; atol=1e-4)

function numerical_diff_symmetric(func, ssinputs_dict::Dict{String, Any}, shock_dict::Dict{String, Any}, h::Float64=1E-4)
    """Same as numerical_diff, but differentiate numerically using central (symmetric) difference, i.e.

    f'(xss)*shock = (f(xss + h*shock) - f(xss - h*shock))/(2*h)
    """

    # response to small shock in each direction
    shocked_inputs_up = Dict(k => ssinputs_dict[k] + h * get(shock_dict, k, 0.0) for k in keys(ssinputs_dict))
    y_up_list = make_tuple(func(shocked_inputs_up...))

    shocked_inputs_down = Dict(k => ssinputs_dict[k] - h * get(shock_dict, k, 0.0) for k in keys(ssinputs_dict))
    y_down_list = make_tuple(func(shocked_inputs_down...))

    # scale responses bak up, dividing by h
    dy_list = [(y_up - y_down) / (2*h)for (y_up, y_down) in zip(y_up_list, y_down_list)]

    return dy_list
end

# test case for numerical_diff_symmetric let's see if this works not confident

# using Test

# Define a simple test function to differentiate
# function test_func(x, y)
#     return x^2 + y^2
# end

# Define the steady-state inputs and shocks
# ssinputs_dict = Dict("x" => 1.0, "y" => 2.0)
# shock_dict = Dict("x" => 0.1, "y" => 0.2)

# Expected derivatives (analytically: df/dx = 2x, df/dy = 2y)
# expected_diff = [2.0 * ssinputs_dict["x"] * shock_dict["x"], 2.0 * ssinputs_dict["y"] * shock_dict["y"]]

# Testing numerical_diff_symmetric
# dy_list_symmetric = numerical_diff_symmetric(test_func, ssinputs_dict, shock_dict)
# @test isapprox(dy_list_symmetric, expected_diff; atol=1e-4)
